# -*- coding: utf-8 -*-
# @Time : 2024-11-2024/11/30
# @File : app.py

import os
import uuid
import time
from flask import Flask, request, jsonify, send_file
from werkzeug.utils import secure_filename
import cv2
import numpy as np # Added missing import for np
import json # Added missing import for json

# 导入我们的姿态分析API
from api import PoseAnalysisAPI

# 初始化Flask应用
app = Flask(__name__)

# 配置上传文件相关参数
UPLOAD_FOLDER = './uploads'
RESULTS_FOLDER = './results'
ALLOWED_EXTENSIONS = {'mp4', 'avi', 'mov', 'mkv'}
MAX_CONTENT_LENGTH = 500 * 1024 * 1024  # 500MB

app.config['UPLOAD_FOLDER'] = UPLOAD_FOLDER
app.config['RESULTS_FOLDER'] = RESULTS_FOLDER
app.config['MAX_CONTENT_LENGTH'] = MAX_CONTENT_LENGTH

# 确保文件夹存在
os.makedirs(UPLOAD_FOLDER, exist_ok=True)
os.makedirs(RESULTS_FOLDER, exist_ok=True)

# 初始化姿态分析API
pose_api = PoseAnalysisAPI(model_path="./weights/yolo11x-pose.pt")

def allowed_file(filename):
    """检查文件类型是否允许"""
    return '.' in filename and \
           filename.rsplit('.', 1)[1].lower() in ALLOWED_EXTENSIONS

@app.route('/')
def index():
    return jsonify({
        "code": 200,
        "msg": "姿态分析API服务正在运行",
        "data": {
            "endpoints": {
                "analyze_video": "/api/analyze_video",
                "analyze_frame": "/api/analyze_frame",
                "get_result": "/api/result/<result_id>",
                "get_result_video": "/api/result_video/<result_id>"
            }
        }
    })

@app.route('/api/analyze_video', methods=['POST'])
def analyze_video():
    """
    分析上传的视频文件
    
    请求参数:
    - video: 视频文件
    
    返回:
    - JSON格式的分析结果
    """
    # 检查是否有文件上传
    if 'video' not in request.files:
        return jsonify({
            "code": 400, 
            "msg": "请求中没有视频文件", 
            "data": None
        })
    
    file = request.files['video']
    
    # 检查文件名是否为空
    if file.filename == '':
        return jsonify({
            "code": 400, 
            "msg": "未选择文件", 
            "data": None
        })
    
    # 检查文件类型
    if not allowed_file(file.filename):
        return jsonify({
            "code": 400, 
            "msg": f"不支持的文件类型，允许的类型：{', '.join(ALLOWED_EXTENSIONS)}", 
            "data": None
        })
    
    try:
        # 生成唯一文件名
        unique_id = str(uuid.uuid4())
        original_filename = secure_filename(file.filename)
        filename_base, filename_ext = os.path.splitext(original_filename)
        unique_filename = f"{filename_base}_{unique_id}{filename_ext}"
        filepath = os.path.join(app.config['UPLOAD_FOLDER'], unique_filename)
        
        # 保存上传的文件
        file.save(filepath)
        
        # 分析视频
        start_time = time.time()
        result = pose_api.analyze_video(filepath, save_result_video=True)
        end_time = time.time()
        
        # 保存结果到JSON文件
        result_filename = f"{filename_base}_{unique_id}_result.json"
        result_filepath = os.path.join(app.config['RESULTS_FOLDER'], result_filename)
        save_result = pose_api.save_result_to_json(result, result_filepath)
        
        # 返回结果
        return jsonify({
            "code": 200,
            "msg": "视频分析完成",
            "data": {
                "resultId": unique_id,
                "originalFilename": original_filename,
                "processingTime": f"{end_time - start_time:.2f}秒",
                "resultFile": result_filename,
                "summary": {
                    "framesProcessed": result["data"]["videoInfo"]["processedFrames"],
                    "totalFrames": result["data"]["videoInfo"]["totalFrames"],
                    "fps": result["data"]["videoInfo"]["fps"]
                }
            }
        })
        
    except Exception as e:
        # 发生错误时返回错误信息
        return jsonify({
            "code": 500,
            "msg": f"处理视频时发生错误: {str(e)}",
            "data": None
        })

@app.route('/api/analyze_frame', methods=['POST'])
def analyze_frame():
    """
    分析单帧图像
    
    请求参数:
    - image: 图像文件
    
    返回:
    - JSON格式的分析结果
    """
    # 检查是否有文件上传
    if 'image' not in request.files:
        return jsonify({
            "code": 400, 
            "msg": "请求中没有图像文件", 
            "data": None
        })
    
    file = request.files['image']
    
    # 检查文件名是否为空
    if file.filename == '':
        return jsonify({
            "code": 400, 
            "msg": "未选择文件", 
            "data": None
        })
    
    try:
        # 读取图像数据
        file_bytes = file.read()
        nparr = np.frombuffer(file_bytes, np.uint8)
        image = cv2.imdecode(nparr, cv2.IMREAD_COLOR)
        
        # 分析图像
        result = pose_api.analyze_frame(image)
        
        # 返回结果
        return jsonify(result)
        
    except Exception as e:
        # 发生错误时返回错误信息
        return jsonify({
            "code": 500,
            "msg": f"处理图像时发生错误: {str(e)}",
            "data": None
        })

@app.route('/api/result_video/<result_id>', methods=['GET'])
def get_result_video(result_id):
    """
    获取分析结果视频文件
    
    请求参数:
    - result_id: 结果ID
    
    返回:
    - 结果视频文件
    """
    try:
        # 查找对应的结果视频文件，优先查找H.264编码的视频
        h264_files = [f for f in os.listdir(app.config['RESULTS_FOLDER']) 
                     if f.find(result_id) != -1 and f.endswith('_result_h264.mp4')]
        
        if h264_files:
            result_files = h264_files
        else:
            result_files = [f for f in os.listdir(app.config['RESULTS_FOLDER']) 
                           if f.find(result_id) != -1 and f.endswith('_result.mp4')]
        
        if not result_files:
            return jsonify({
                "code": 404,
                "message": f"未找到ID为 {result_id} 的结果视频文件",
                "data": None
            })
        
        result_file = result_files[0]
        result_path = os.path.join(app.config['RESULTS_FOLDER'], result_file)
        
        # 返回结果视频文件
        return send_file(result_path, mimetype='video/mp4', as_attachment=True)
        
    except Exception as e:
        # 发生错误时返回错误信息
        return jsonify({
            "code": 500,
            "message": f"获取结果视频时发生错误: {str(e)}",
            "data": None
        })

@app.route('/api/result/<result_id>', methods=['GET'])
def get_result(result_id):
    """
    获取分析结果文件
    
    请求参数:
    - result_id: 结果ID
    
    返回:
    - JSON格式的分析结果文件
    """
    try:
        # 查找对应的结果文件
        result_files = [f for f in os.listdir(app.config['RESULTS_FOLDER']) 
                       if f.find(result_id) != -1 and f.endswith('_result.json')]
        
        if not result_files:
            return jsonify({
                "code": 404,
                "msg": f"未找到ID为 {result_id} 的结果文件",
                "data": None
            })
        
        result_file = result_files[0]
        result_path = os.path.join(app.config['RESULTS_FOLDER'], result_file)
        
        # 读取结果文件
        with open(result_path, 'r', encoding='utf-8') as f:
            result_data = json.load(f)
            
        # 移除所有 keypoints 字段
        if result_data.get('data') and result_data['data'].get('frames'):
            for frame in result_data['data']['frames']:
                if frame.get('data'):
                    for person in frame['data']:
                        if 'keypoints' in person:
                            del person['keypoints']
        
        # 创建一个新的临时文件以保存修改后的结果
        temp_result_path = os.path.join(app.config['RESULTS_FOLDER'], f"temp_{result_file}")
        with open(temp_result_path, 'w', encoding='utf-8') as f:
            json.dump(result_data, f, ensure_ascii=False, indent=2)
        
        # 返回修改后的结果文件
        return send_file(temp_result_path, mimetype='application/json', as_attachment=True)
        
    except Exception as e:
        # 发生错误时返回错误信息
        return jsonify({
            "code": 500,
            "msg": f"获取结果时发生错误: {str(e)}",
            "data": None
        })

if __name__ == '__main__':
    app.run(host='0.0.0.0', port=5000, debug=True) 